"""
Memory Cleanup Nodes for ComfyUI

Provides manual and automatic memory cleanup functionality.
"""

import time
from ..utils.memory_utils import cleanup_memory, optimize_vram, get_memory_info, format_bytes, global_memory_manager


class MemoryCleanupNode:
    """
    Manual Memory Cleanup Node - Performs immediate memory cleanup
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "trigger": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 1000000,
                    "step": 1,
                    "display": "number"
                }),
                "aggressive_cleanup": ("BOOLEAN", {"default": False}),
                "include_vram": ("BOOLEAN", {"default": True}),
            },
        }
    
    RETURN_TYPES = (
        "STRING",      # Cleanup report
        "INT",         # Objects collected
        "INT",         # Memory freed (bytes)
        "BOOLEAN",     # Success status
    )
    
    RETURN_NAMES = (
        "cleanup_report",
        "objects_collected",
        "memory_freed",
        "success",
    )
    
    FUNCTION = "cleanup_memory"
    CATEGORY = "Memory Management"
    OUTPUT_NODE = True
    
    def cleanup_memory(self, trigger, aggressive_cleanup=False, include_vram=True):
        """Perform memory cleanup and return results"""
        
        try:
            # Get initial memory state
            initial_memory = get_memory_info()
            
            # Perform cleanup
            cleanup_result = cleanup_memory(aggressive=aggressive_cleanup)
            
            # Optimize VRAM if requested
            vram_result = {}
            if include_vram:
                vram_result = optimize_vram()
            
            # Get final memory state
            final_memory = get_memory_info()
            
            # Create report
            report = self._create_cleanup_report(
                initial_memory, 
                final_memory, 
                cleanup_result, 
                vram_result,
                aggressive_cleanup,
                include_vram
            )
            
            return (
                report,
                cleanup_result.get('objects_collected', 0),
                cleanup_result.get('memory_freed', 0),
                True
            )
            
        except Exception as e:
            error_report = f"Memory cleanup failed: {str(e)}"
            return (
                error_report,
                0,
                0,
                False
            )
    
    def _create_cleanup_report(self, initial_memory, final_memory, cleanup_result, vram_result, aggressive, include_vram):
        """Create detailed cleanup report"""
        
        report_lines = ["=== MEMORY CLEANUP REPORT ===\n"]
        
        # Cleanup mode
        mode = "Aggressive" if aggressive else "Standard"
        report_lines.append(f"Cleanup Mode: {mode}")
        report_lines.append(f"VRAM Cleanup: {'Enabled' if include_vram else 'Disabled'}")
        
        # Objects collected
        objects_collected = cleanup_result.get('objects_collected', 0)
        report_lines.append(f"\nObjects Collected: {objects_collected}")
        
        # Memory changes
        initial_rss = initial_memory.get('process_rss', 0)
        final_rss = final_memory.get('process_rss', 0)
        memory_freed = initial_rss - final_rss
        
        report_lines.append(f"\nRAM CHANGES:")
        report_lines.append(f"  Before: {format_bytes(initial_rss)}")
        report_lines.append(f"  After: {format_bytes(final_rss)}")
        report_lines.append(f"  Freed: {format_bytes(memory_freed)}")
        
        # System memory changes
        initial_system = initial_memory.get('system_percent', 0)
        final_system = final_memory.get('system_percent', 0)
        system_improvement = initial_system - final_system
        
        report_lines.append(f"\nSYSTEM MEMORY:")
        report_lines.append(f"  Before: {initial_system:.1f}%")
        report_lines.append(f"  After: {final_system:.1f}%")
        report_lines.append(f"  Improvement: {system_improvement:.1f}%")
        
        # VRAM changes if applicable
        if include_vram and vram_result and 'status' not in vram_result:
            report_lines.append(f"\nVRAM CLEANUP:")
            for device_key, device_info in vram_result.items():
                if isinstance(device_info, dict):
                    freed_allocated = device_info.get('freed_allocated', 0)
                    freed_reserved = device_info.get('freed_reserved', 0)
                    report_lines.append(f"  {device_key.upper()}:")
                    report_lines.append(f"    Allocated Freed: {format_bytes(freed_allocated)}")
                    report_lines.append(f"    Reserved Freed: {format_bytes(freed_reserved)}")
        
        return "\n".join(report_lines)


class AutoMemoryCleanupNode:
    """
    Automatic Memory Cleanup Node - Configures automatic memory management
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "enable_monitoring": ("BOOLEAN", {"default": True}),
                "warning_threshold": ("FLOAT", {
                    "default": 80.0,
                    "min": 60.0,
                    "max": 95.0,
                    "step": 5.0,
                    "display": "slider"
                }),
                "critical_threshold": ("FLOAT", {
                    "default": 90.0,
                    "min": 70.0,
                    "max": 98.0,
                    "step": 2.0,
                    "display": "slider"
                }),
                "check_interval": ("FLOAT", {
                    "default": 30.0,
                    "min": 5.0,
                    "max": 300.0,
                    "step": 5.0,
                    "display": "slider"
                }),
            },
        }
    
    RETURN_TYPES = (
        "STRING",      # Status report
        "BOOLEAN",     # Is monitoring active
        "FLOAT",       # Current memory percentage
        "STRING",      # Last action taken
    )
    
    RETURN_NAMES = (
        "status_report",
        "monitoring_active",
        "current_memory_percent",
        "last_action",
    )
    
    FUNCTION = "configure_auto_cleanup"
    CATEGORY = "Memory Management"
    OUTPUT_NODE = True
    
    def configure_auto_cleanup(self, enable_monitoring, warning_threshold, critical_threshold, check_interval):
        """Configure automatic memory cleanup"""
        
        try:
            # Update global memory manager settings
            global_memory_manager.warning_threshold = warning_threshold
            global_memory_manager.critical_threshold = critical_threshold
            global_memory_manager.check_interval = check_interval
            
            # Start or stop monitoring based on setting
            if enable_monitoring and not global_memory_manager.running:
                global_memory_manager.start_monitoring()
                action = "Started automatic monitoring"
            elif not enable_monitoring and global_memory_manager.running:
                global_memory_manager.stop_monitoring()
                action = "Stopped automatic monitoring"
            else:
                action = "Configuration updated"
            
            # Get current memory state
            current_memory = get_memory_info()
            memory_percent = current_memory.get('system_percent', 0)
            
            # Create status report
            status_report = self._create_status_report(
                enable_monitoring,
                warning_threshold,
                critical_threshold,
                check_interval,
                memory_percent,
                action
            )
            
            return (
                status_report,
                global_memory_manager.running,
                memory_percent,
                action
            )
            
        except Exception as e:
            error_report = f"Auto-cleanup configuration failed: {str(e)}"
            return (
                error_report,
                False,
                0.0,
                "Error occurred"
            )
    
    def _create_status_report(self, enabled, warning_threshold, critical_threshold, check_interval, memory_percent, action):
        """Create status report for auto-cleanup"""
        
        report_lines = ["=== AUTO MEMORY CLEANUP STATUS ===\n"]
        
        # Configuration
        report_lines.append(f"Monitoring: {'ENABLED' if enabled else 'DISABLED'}")
        report_lines.append(f"Warning Threshold: {warning_threshold:.1f}%")
        report_lines.append(f"Critical Threshold: {critical_threshold:.1f}%")
        report_lines.append(f"Check Interval: {check_interval:.1f} seconds")
        
        # Current status
        report_lines.append(f"\nCURRENT STATUS:")
        report_lines.append(f"  Memory Usage: {memory_percent:.1f}%")
        report_lines.append(f"  Last Action: {action}")
        
        # Status indicators
        if memory_percent >= critical_threshold:
            status = "🔴 CRITICAL - Automatic cleanup active"
        elif memory_percent >= warning_threshold:
            status = "🟡 WARNING - Cleanup will trigger soon"
        else:
            status = "🟢 NORMAL - Memory usage within limits"
        
        report_lines.append(f"  Status: {status}")
        
        # Recommendations
        report_lines.append(f"\nRECOMMENDATIONS:")
        if not enabled:
            report_lines.append("  • Enable monitoring for automatic memory management")
        elif memory_percent >= critical_threshold:
            report_lines.append("  • Consider reducing memory usage manually")
            report_lines.append("  • Close unnecessary applications")
        elif memory_percent >= warning_threshold:
            report_lines.append("  • Monitor memory usage closely")
            report_lines.append("  • Prepare for potential cleanup")
        else:
            report_lines.append("  • Current configuration is optimal")
        
        return "\n".join(report_lines) 